import java.io.BufferedOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

public class Huff implements ITreeMaker, IHuffEncoder, IHuffModel, IHuffHeader {
    private final Map<Integer, String> encodingMap;
    private Map<Integer, Integer> countMap;
    private HuffTree huffTree;
    private final ICharCounter myCounter;

    public Huff() {
        encodingMap = new HashMap<>();
        myCounter = new CharCounter();
    }

    @Override
    public HuffTree makeHuffTree(InputStream stream) throws IOException {
        // count characters
        myCounter.countAll(stream);
        countMap = myCounter.getTable();

        // create priority queue for tree construction
        PriorityQueue<HuffTree> pq = new PriorityQueue<>();

        // add leaf nodes
        for (Map.Entry<Integer, Integer> entry : countMap.entrySet()) {
            pq.add(new HuffTree(entry.getKey(), entry.getValue()));
        }

        // add EOF
        pq.add(new HuffTree(PSEUDO_EOF, 1));

        // build tree
        while (pq.size() > 1) {
            HuffTree left = pq.poll();
            HuffTree right = pq.poll();
            if (right != null) {
                pq.add(new HuffTree(left.root(), right.root(), left.weight() + right.weight()));
            }
        }

        huffTree = pq.poll();
        return huffTree;
    }

    @Override
    public Map<Integer, String> makeTable() {
        encodingMap.clear();
        if (huffTree != null) {
            buildEncodingMap(huffTree.root(), "");
        }
        return encodingMap;
    }

    private void buildEncodingMap(IHuffBaseNode node, String path) {
        if (node == null) {
            return;
        }

        if (node.isLeaf()) {
            HuffLeafNode leaf = (HuffLeafNode) node;
            encodingMap.put(leaf.element(), path);
        } else {
            HuffInternalNode internal = (HuffInternalNode) node;
            buildEncodingMap(internal.left(), path + "0");
            buildEncodingMap(internal.right(), path + "1");
        }
    }

    @Override
    public String getCode(int i) {
        if (!encodingMap.containsKey(i)) {
            throw new IllegalStateException("No encoding found for character " + i);
        }
        return encodingMap.get(i);
    }

    @Override
    public Map<Integer, Integer> showCounts() {
        if (countMap == null) {
            throw new IllegalStateException("Counts not available. Call makeHuffTree first.");
        }
        return new HashMap<>(countMap);
    }

    @Override
    public int headerSize() {
        // magic num (32 bits) + tree size
        return BITS_PER_INT + treeSize(huffTree.root());
    }

    private int treeSize(IHuffBaseNode node) {
        if (node == null) {
            return 0;
        }
        if (node.isLeaf()) {
            return 1 + (BITS_PER_WORD + 1); // 1 bit flag + 9 bits value
        }
        return 1 + treeSize(((HuffInternalNode) node).left())
                + treeSize(((HuffInternalNode) node).right());
    }

    @Override
    public int writeHeader(BitOutputStream out) {
        try {
            int bits = 0;
            out.write(BITS_PER_INT, MAGIC_NUMBER);
            bits += BITS_PER_INT;
            bits += writeTree(huffTree.root(), out);
            return bits;
        } catch (IOException e) {
            throw new RuntimeException("Error writing header", e);
        }
    }

    private int writeTree(IHuffBaseNode node, BitOutputStream out) throws IOException {
        int bits = 0;
        if (node.isLeaf()) {
            out.write(1, 1); // mark leaf node
            out.write(BITS_PER_WORD + 1, ((HuffLeafNode) node).element());
            bits += 1 + (BITS_PER_WORD + 1); // total bits count
        } else {
            out.write(1, 0); // mark internal node
            bits += 1;
            HuffInternalNode internal = (HuffInternalNode) node;
            bits += writeTree(internal.left(), out);
            bits += writeTree(internal.right(), out);
        }
        return bits;
    }

    @Override
    public HuffTree readHeader(BitInputStream in) throws IOException {
        int magic = in.read(BITS_PER_INT);
        if (magic != MAGIC_NUMBER) {
            throw new IOException("Bad magic number");
        }

        // read and build tree directly
        IHuffBaseNode root = readTree(in);
        if (root.isLeaf()) {
            HuffLeafNode leaf = (HuffLeafNode) root;
            return new HuffTree(leaf.element(), 0);
        } else {
            HuffInternalNode internal = (HuffInternalNode) root;
            return new HuffTree(internal.left(), internal.right(), 0);
        }
    }

    private IHuffBaseNode readTree(BitInputStream in) throws IOException {
        int bit = in.read(1);

        if (bit == -1) {
            throw new IOException("Unexpected end of file");
        }

        if (bit == 1) { // leaf node
            int value = in.read(BITS_PER_WORD + 1);
            if (value == -1) {
                throw new IOException("Unexpected end of file");
            }
            return new HuffLeafNode(value, 0);
        } else { // internal node
            IHuffBaseNode left = readTree(in);
            IHuffBaseNode right = readTree(in);
            return new HuffInternalNode(left, right, 0);
        }
    }

    @Override
    public int write(String inFile, String outFile, boolean force) {
        try {
            // first pass: count frequencies and build tree
            try (InputStream in = Files.newInputStream(Paths.get(inFile))) {
                makeHuffTree(in);
                makeTable();
            }

            // calculate compressed size in bits
            int compressedBits = calculateCompressedSize();
            File inputFile = new File(inFile);
            long originalBits = inputFile.length() * 8; // Convert bytes to bits

            // don't compress if it would make the file larger, unless forced
            if (!force && compressedBits >= originalBits) {
                return compressedBits;
            }

            // second pass: write compressed file
            try (InputStream in = Files.newInputStream(Paths.get(inFile));
                    BitOutputStream out = new BitOutputStream(outFile)) {

                // write header
                writeHeader(out);

                // write encoded data
                int ch;
                while ((ch = in.read()) != -1) {
                    String code = getCode(ch);
                    for (char bit : code.toCharArray()) {
                        out.write(1, bit - '0');
                    }
                }

                // write PSEUDO_EOF
                String eofCode = getCode(PSEUDO_EOF);
                for (char bit : eofCode.toCharArray()) {
                    out.write(1, bit - '0');
                }

                return compressedBits;
            }
        } catch (IOException e) {
            throw new RuntimeException("Error during compression", e);
        }
    }

    private int calculateCompressedSize() {
        int size = headerSize(); // include header size
        for (Map.Entry<Integer, Integer> entry : countMap.entrySet()) {
            String code = encodingMap.get(entry.getKey());
            if (code != null) {
                size += code.length() * entry.getValue();
            }
        }
        // PSEUDO_EOF
        String eofCode = encodingMap.get(PSEUDO_EOF);
        if (eofCode != null) {
            size += eofCode.length();
        }
        return size;
    }

    @Override
    public int uncompress(String inFile, String outFile) {
        try (BitInputStream in = new BitInputStream(inFile);
                BufferedOutputStream out =
                        new BufferedOutputStream(Files.newOutputStream(Paths.get(outFile)))) {

            huffTree = readHeader(in); // read tree after magic num
            if (huffTree == null || huffTree.root() == null) {
                throw new IOException("bad tree structure");
            }

            IHuffBaseNode current = huffTree.root();
            int totalChars = 0;

            while (true) {
                if (current.isLeaf()) {
                    HuffLeafNode leaf = (HuffLeafNode) current;
                    int value = leaf.element();

                    if (value == PSEUDO_EOF) {
                        break; // stop when meet pseudo eof
                    }

                    out.write(value);
                    totalChars++;
                    current = huffTree.root();
                } else {
                    int bit = in.read(1);
                    if (bit == -1) {
                        throw new IOException("Unexpected end of file while reading bits");
                    }

                    HuffInternalNode node = (HuffInternalNode) current;
                    current = (bit == 0) ? node.left() : node.right();

                    if (current == null) {
                        throw new IOException("Corrupted Huffman tree structure");
                    }
                }
            }

            out.flush();
            return totalChars * 8;

        } catch (IOException e) {
            return 0;
        }
    }
}
